"""
增强版空调负荷预测Transformer模型
基于查表嵌入的时间特征处理和改进的架构设计
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from typing import Optional, Tuple, Dict
import datetime
from dataclasses import dataclass
import math


@dataclass
class EnhancedHVACConfig:
    """增强版空调负荷预测配置"""
    # 模型参数
    d_model: int = 256          # 模型维度
    n_heads: int = 8            # 注意力头数
    n_layers: int = 6           # Transformer层数
    d_ff: int = 1024           # 前馈网络维度
    dropout: float = 0.1        # Dropout率
    
    # 时间参数
    seq_len: int = 168          # 输入序列长度（一周）
    pred_len: int = 24          # 预测长度（一天）
    
    # 特征参数
    n_env_features: int = 2     # 环境特征数量（温度、湿度）
    n_cyclical_features: int = 6 # 周期性特征数量
    n_statistical_features: int = 8 # 统计特征数量
    
    # 嵌入参数
    embedding_dims: Dict[str, int] = None  # 各分类特征的嵌入维度
    
    # 设备配置
    device: str = "cuda" if torch.cuda.is_available() else "cpu"
    
    def __post_init__(self):
        if self.embedding_dims is None:
            self.embedding_dims = {
                'hour': 32,        # 小时嵌入维度
                'dayofweek': 16,   # 星期几嵌入维度
                'month': 16,       # 月份嵌入维度
                'season': 8,       # 季节嵌入维度
                'is_holiday': 8,   # 节假日嵌入维度
                'is_weekend': 8,   # 周末嵌入维度
                'work_period': 16, # 工作时段嵌入维度
            }


class PositionalEncoding(nn.Module):
    """位置编码"""
    
    def __init__(self, d_model: int, max_len: int = 5000):
        super().__init__()
        
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * 
                           (-math.log(10000.0) / d_model))
        
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        pe = pe.unsqueeze(0).transpose(0, 1)
        
        self.register_buffer('pe', pe)
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return x + self.pe[:x.size(0), :]


class CategoricalEmbedding(nn.Module):
    """分类特征嵌入层"""
    
    def __init__(self, config: EnhancedHVACConfig):
        super().__init__()
        self.config = config
        
        # 创建各种分类特征的嵌入层
        self.embeddings = nn.ModuleDict({
            'hour': nn.Embedding(24, config.embedding_dims['hour']),
            'dayofweek': nn.Embedding(7, config.embedding_dims['dayofweek']),
            'month': nn.Embedding(12, config.embedding_dims['month']),
            'season': nn.Embedding(4, config.embedding_dims['season']),
            'is_holiday': nn.Embedding(2, config.embedding_dims['is_holiday']),
            'is_weekend': nn.Embedding(2, config.embedding_dims['is_weekend']),
            'work_period': nn.Embedding(4, config.embedding_dims['work_period']),
        })
        
        # 计算总嵌入维度
        self.total_embed_dim = sum(config.embedding_dims.values())
        
        # 嵌入融合层
        self.embed_fusion = nn.Linear(self.total_embed_dim, config.d_model // 4)
        self.dropout = nn.Dropout(config.dropout)
        
    def forward(self, categorical_features: Dict[str, torch.Tensor]) -> torch.Tensor:
        """
        Args:
            categorical_features: 分类特征字典，每个值为 [batch_size, seq_len]
        Returns:
            embedded_features: [batch_size, seq_len, d_model//4]
        """
        embedded_list = []
        
        for feature_name, embedding_layer in self.embeddings.items():
            if feature_name in categorical_features:
                # [batch_size, seq_len] -> [batch_size, seq_len, embed_dim]
                embedded = embedding_layer(categorical_features[feature_name])
                embedded_list.append(embedded)
        
        # 拼接所有嵌入特征
        if embedded_list:
            # [batch_size, seq_len, total_embed_dim]
            concatenated = torch.cat(embedded_list, dim=-1)
            # 融合到目标维度
            fused = self.embed_fusion(concatenated)
            return self.dropout(fused)
        else:
            # 如果没有分类特征，返回零张量
            batch_size, seq_len = list(categorical_features.values())[0].shape[:2]
            return torch.zeros(batch_size, seq_len, self.config.d_model // 4, 
                             device=next(self.parameters()).device)


class MultiFeatureEncoder(nn.Module):
    """多特征编码器"""
    
    def __init__(self, config: EnhancedHVACConfig):
        super().__init__()
        self.config = config
        
        # 负荷值投影
        self.load_projection = nn.Linear(1, config.d_model // 4)
        
        # 环境特征投影
        self.env_projection = nn.Linear(config.n_env_features, config.d_model // 4)
        
        # 周期性特征投影
        self.cyclical_projection = nn.Linear(config.n_cyclical_features, config.d_model // 8)
        
        # 统计特征投影
        self.statistical_projection = nn.Linear(config.n_statistical_features, config.d_model // 8)
        
        # 分类特征嵌入
        self.categorical_embedding = CategoricalEmbedding(config)
        
        # 动态特征融合层（会根据实际输入特征动态创建）
        self.feature_fusion = None
        self.fusion_input_dim = None
        
    def forward(self, features: Dict[str, torch.Tensor]) -> torch.Tensor:
        """
        Args:
            features: 特征字典
        Returns:
            encoded_features: [batch_size, seq_len, d_model]
        """
        feature_list = []
        
        # 处理连续特征
        if 'load' in features:
            load_embed = self.load_projection(features['load'])
            feature_list.append(load_embed)
        
        if 'env' in features:
            env_embed = self.env_projection(features['env'])
            feature_list.append(env_embed)
        
        if 'cyclical_features' in features:
            cyclical_embed = self.cyclical_projection(features['cyclical_features'])
            feature_list.append(cyclical_embed)
        
        if 'statistical_features' in features:
            statistical_embed = self.statistical_projection(features['statistical_features'])
            feature_list.append(statistical_embed)
        
        # 处理分类特征
        categorical_features = {
            k: v for k, v in features.items() 
            if k in ['hour', 'dayofweek', 'month', 'season', 'is_holiday', 'is_weekend', 'work_period']
        }
        if categorical_features:
            categorical_embed = self.categorical_embedding(categorical_features)
            feature_list.append(categorical_embed)
        
        # 拼接所有特征
        if feature_list:
            concatenated = torch.cat(feature_list, dim=-1)
            
            # 动态创建融合层（仅在第一次调用时）
            if self.feature_fusion is None:
                actual_input_dim = concatenated.shape[-1]
                self.fusion_input_dim = actual_input_dim
                self.feature_fusion = nn.Sequential(
                    nn.Linear(actual_input_dim, self.config.d_model),
                    nn.LayerNorm(self.config.d_model),
                    nn.ReLU(),
                    nn.Dropout(self.config.dropout)
                ).to(concatenated.device)
                print(f"🔧 动态创建特征融合层: {actual_input_dim} -> {self.config.d_model}")
            
            return self.feature_fusion(concatenated)
        else:
            raise ValueError("没有提供有效的输入特征")


class TransformerEncoderLayer(nn.Module):
    """Transformer编码器层"""
    
    def __init__(self, config: EnhancedHVACConfig):
        super().__init__()
        self.self_attention = nn.MultiheadAttention(
            config.d_model, config.n_heads, config.dropout, batch_first=True
        )
        self.feed_forward = nn.Sequential(
            nn.Linear(config.d_model, config.d_ff),
            nn.ReLU(),
            nn.Dropout(config.dropout),
            nn.Linear(config.d_ff, config.d_model)
        )
        self.norm1 = nn.LayerNorm(config.d_model)
        self.norm2 = nn.LayerNorm(config.d_model)
        self.dropout = nn.Dropout(config.dropout)
        
    def forward(self, x: torch.Tensor, mask: Optional[torch.Tensor] = None) -> torch.Tensor:
        # 自注意力
        attn_output, _ = self.self_attention(x, x, x, attn_mask=mask)
        x = self.norm1(x + self.dropout(attn_output))
        
        # 前馈网络
        ff_output = self.feed_forward(x)
        x = self.norm2(x + self.dropout(ff_output))
        
        return x


class EnhancedHVACTransformer(nn.Module):
    """增强版空调负荷预测Transformer"""
    
    def __init__(self, config: EnhancedHVACConfig):
        super().__init__()
        self.config = config
        
        # 多特征编码器
        self.feature_encoder = MultiFeatureEncoder(config)
        
        # 位置编码
        self.positional_encoding = PositionalEncoding(config.d_model, config.seq_len + config.pred_len)
        
        # Transformer编码器层
        self.transformer_layers = nn.ModuleList([
            TransformerEncoderLayer(config) for _ in range(config.n_layers)
        ])
        
        # 预测头
        self.prediction_head = nn.Sequential(
            nn.Linear(config.d_model, config.d_model // 2),
            nn.ReLU(),
            nn.Dropout(config.dropout),
            nn.Linear(config.d_model // 2, config.d_model // 4),
            nn.ReLU(),
            nn.Dropout(config.dropout),
            nn.Linear(config.d_model // 4, 1)
        )
        
        # 自回归预测的初始状态
        self.prediction_start_token = nn.Parameter(torch.randn(1, 1, config.d_model))
        
    def create_causal_mask(self, seq_len: int) -> torch.Tensor:
        """创建因果掩码"""
        mask = torch.triu(torch.ones(seq_len, seq_len), diagonal=1)
        mask = mask.masked_fill(mask == 1, float('-inf'))
        return mask.to(next(self.parameters()).device)
    
    def forward(self, features: Dict[str, torch.Tensor], 
                pred_len: Optional[int] = None) -> torch.Tensor:
        """
        Args:
            features: 输入特征字典
            pred_len: 预测长度
        Returns:
            predictions: [batch_size, pred_len, 1]
        """
        if pred_len is None:
            pred_len = self.config.pred_len
        
        batch_size, seq_len = features['load'].shape[:2]
        
        # 特征编码
        encoded_features = self.feature_encoder(features)  # [batch_size, seq_len, d_model]
        
        # 添加位置编码
        encoded_features = self.positional_encoding(encoded_features.transpose(0, 1)).transpose(0, 1)
        
        # 通过Transformer编码器
        mask = self.create_causal_mask(seq_len)
        hidden_states = encoded_features
        
        for layer in self.transformer_layers:
            hidden_states = layer(hidden_states, mask)
        
        # 自回归预测
        predictions = []
        current_hidden = hidden_states[:, -1:, :]  # 最后一个时间步的隐藏状态
        
        for step in range(pred_len):
            # 预测当前步
            pred_output = self.prediction_head(current_hidden)  # [batch_size, 1, 1]
            predictions.append(pred_output)
            
            # 为下一步准备输入（简化版本，实际应用中需要考虑新的环境特征）
            if step < pred_len - 1:
                # 创建下一个时间步的特征（这里使用简化的方法）
                next_features = current_hidden  # 实际应用中需要结合新的环境数据
                current_hidden = next_features
        
        return torch.cat(predictions, dim=1)  # [batch_size, pred_len, 1]


class RealTimeEnhancedPredictor:
    """增强版实时预测器"""
    
    def __init__(self, model: EnhancedHVACTransformer, config: EnhancedHVACConfig):
        self.model = model
        self.config = config
        self.model.eval()
        
        # 数据缓存
        self.feature_buffer = {}
        self.max_buffer_size = config.seq_len
        
    def update_data(self, load_value: float, env_data: np.ndarray, 
                   timestamp: datetime.datetime):
        """
        更新实时数据
        
        Args:
            load_value: 当前负荷值
            env_data: 环境数据 [temperature, humidity]
            timestamp: 时间戳
        """
        # 初始化缓冲区
        if not self.feature_buffer:
            feature_names = ['load', 'env', 'hour', 'dayofweek', 'month', 'season', 
                           'is_holiday', 'is_weekend', 'work_period']
            for name in feature_names:
                self.feature_buffer[name] = []
        
        # 添加新数据
        self.feature_buffer['load'].append(load_value)
        self.feature_buffer['env'].append(env_data)
        self.feature_buffer['hour'].append(timestamp.hour)
        self.feature_buffer['dayofweek'].append(timestamp.weekday())
        self.feature_buffer['month'].append(timestamp.month - 1)
        self.feature_buffer['season'].append((timestamp.month - 1) // 3)
        
        # 简化的节假日检测（实际应用中需要完整的节假日表）
        date_str = timestamp.strftime('%Y-%m-%d')
        self.feature_buffer['is_holiday'].append(0)  # 简化为无节假日
        
        self.feature_buffer['is_weekend'].append(1 if timestamp.weekday() >= 5 else 0)
        
        # 工作时段
        hour = timestamp.hour
        if 6 <= hour < 12:
            work_period = 1  # 上午
        elif 12 <= hour < 18:
            work_period = 2  # 下午
        elif 18 <= hour < 22:
            work_period = 3  # 晚上
        else:
            work_period = 0  # 夜间
        self.feature_buffer['work_period'].append(work_period)
        
        # 保持缓冲区大小
        for key in self.feature_buffer:
            if len(self.feature_buffer[key]) > self.max_buffer_size:
                self.feature_buffer[key].pop(0)
    
    def predict_remaining_day(self) -> Tuple[np.ndarray, Dict]:
        """预测今日剩余时间的负荷"""
        if len(self.feature_buffer['load']) < self.config.seq_len:
            raise ValueError(f"需要至少 {self.config.seq_len} 个历史数据点")
        
        # 准备输入特征
        features = {}
        
        # 连续特征
        features['load'] = torch.tensor(
            self.feature_buffer['load'][-self.config.seq_len:], 
            dtype=torch.float32
        ).unsqueeze(0).unsqueeze(-1)  # [1, seq_len, 1]
        
        features['env'] = torch.tensor(
            self.feature_buffer['env'][-self.config.seq_len:], 
            dtype=torch.float32
        ).unsqueeze(0)  # [1, seq_len, 2]
        
        # 分类特征
        categorical_names = ['hour', 'dayofweek', 'month', 'season', 
                           'is_holiday', 'is_weekend', 'work_period']
        for name in categorical_names:
            features[name] = torch.tensor(
                self.feature_buffer[name][-self.config.seq_len:], 
                dtype=torch.long
            ).unsqueeze(0)  # [1, seq_len]
        
        # 移动到设备
        device = next(self.model.parameters()).device
        for key in features:
            features[key] = features[key].to(device)
        
        # 计算今日剩余小时数
        current_hour = datetime.datetime.now().hour
        remaining_hours = 24 - current_hour
        
        # 进行预测
        with torch.no_grad():
            predictions = self.model(features, remaining_hours)
            predictions_np = predictions.cpu().numpy().squeeze()
        
        # 生成时间戳
        current_time = datetime.datetime.now()
        time_stamps = [current_time + datetime.timedelta(hours=i+1) 
                      for i in range(remaining_hours)]
        
        # 计算统计信息
        stats = {
            'mean_load': float(np.mean(predictions_np)),
            'max_load': float(np.max(predictions_np)),
            'min_load': float(np.min(predictions_np)),
            'peak_hour': time_stamps[np.argmax(predictions_np)] if len(predictions_np) > 0 else None,
            'total_consumption': float(np.sum(predictions_np)),
            'remaining_hours': remaining_hours
        }
        
        return predictions_np, stats


def create_enhanced_hvac_model():
    """创建增强版空调负荷预测模型"""
    config = EnhancedHVACConfig(
        d_model=256,
        n_heads=8,
        n_layers=6,
        seq_len=168,  # 一周数据
        pred_len=24,  # 预测一天
        n_env_features=2,  # 温度、湿度
        dropout=0.1
    )
    
    model = EnhancedHVACTransformer(config)
    predictor = RealTimeEnhancedPredictor(model, config)
    
    return model, predictor, config


if __name__ == "__main__":
    # 创建模型
    model, predictor, config = create_enhanced_hvac_model()
    
    print(f"增强模型参数数量: {sum(p.numel() for p in model.parameters()):,}")
    print(f"模型配置: {config}")
    
    # 测试模型前向传播
    batch_size, seq_len = 2, 168
    
    # 创建测试输入
    test_features = {
        'load': torch.randn(batch_size, seq_len, 1),
        'env': torch.randn(batch_size, seq_len, 2),
        'hour': torch.randint(0, 24, (batch_size, seq_len)),
        'dayofweek': torch.randint(0, 7, (batch_size, seq_len)),
        'month': torch.randint(0, 12, (batch_size, seq_len)),
        'season': torch.randint(0, 4, (batch_size, seq_len)),
        'is_holiday': torch.randint(0, 2, (batch_size, seq_len)),
        'is_weekend': torch.randint(0, 2, (batch_size, seq_len)),
        'work_period': torch.randint(0, 4, (batch_size, seq_len)),
    }
    
    print("\n=== 模型测试 ===")
    with torch.no_grad():
        output = model(test_features, pred_len=24)
        print(f"输出形状: {output.shape}")
        print("模型前向传播测试成功！")
